import rospy
from math import *
from controller import Camera
from controller import Robot
from sensor_msgs.msg import Image, CameraInfo

robot=Robot()
robot_name = robot.getName()
timestep = int(robot.getBasicTimeStep())
camera = robot.getDevice('camera')
camera.enable(timestep)

def publish_camera(time):   # publish image to ROS topic
    img_msg = Image()
    img_msg.header.stamp = rospy.Time.from_seconds(time)
    # img_msg.header.frame_id = camera_optical_frame
    img_msg.height = camera.getHeight()
    img_msg.width = camera.getWidth()
    img_msg.encoding = "bgra8"
    img_msg.step = 4 * camera.getWidth()
    img = camera.getImage()
    img_msg.data = img
    pub_cam.publish(img_msg)

def publish_camera_info(time):  # publish camera info to ROS topic
    cam_info = CameraInfo()
    cam_info.header.stamp = rospy.Time.from_seconds(time)
    cam_info.height = camera.getHeight()
    cam_info.width = camera.getWidth()
    f_y = mat_from_fov_and_resolution(
        h_fov_to_v_fov(camera.getFov(), cam_info.height, cam_info.width),
        cam_info.height)
    f_x = mat_from_fov_and_resolution(camera.getFov(), cam_info.width)
    cam_info.K = [f_x, 0, cam_info.width / 2,
                    0, f_y, cam_info.height / 2,
                    0, 0, 1]
    cam_info.P = [f_x, 0, cam_info.width / 2, 0,
                    0, f_y, cam_info.height / 2, 0,
                    0, 0, 1, 0]
    cam_info.header.frame_id = "map"  #need to be modified
    pub_cam_info.publish(cam_info)

def mat_from_fov_and_resolution(fov, res):
    return 0.5*res * (cos((fov / 2)) / sin((fov / 2)))

def h_fov_to_v_fov(h_fov, height, width):
    return 2 * atan2(tan(h_fov * 0.5) * (height),width)

rospy.init_node(robot_name)
pub_cam_info = rospy.Publisher("~camera_info", CameraInfo, queue_size=1, latch=True)
pub_cam = rospy.Publisher("~image_raw", Image, queue_size=1)

if __name__ == "__main__":
    time = 0
    while robot.step(timestep) != -1 and not rospy.is_shutdown():
        publish_camera(time)
        publish_camera_info(time)
        time += 1

